import json
from pathlib import Path

from agents.dqn import DQN


def generate_config():
    config = {
        'use_GPU': True,
        'max_steps': 4e5,
        'eval_freq': 1000,
        'batch_size': 256,
        'buffer_size': 1e5,
        'lr': 1e-4,
        'gamma': 0.99,
        'alpha_per': 0.6,
        'beta_per': 0.4,
        'target_update_freq': 200,
        'tau': 0.005,
        'epsilon_init': 0.5,
        'epsilon_final': 0.1,
        'epsilon_decay_steps': 1e5,
        'grad_clip': 10,
        'use_per': True,
        'use_double': True,
        'use_dueling': True,
        'use_lr_decay': True,
        'use_soft_update': True,
    }
    return config



if __name__ == "__main__":
    env_list = ['CartPole-v1', 'Acrobot-v1', 'MountainCar-v0', 'LunarLander-v3']
    env = env_list[3]
    seed_list = [0, 10]
    config = generate_config()
    results = []
    for seed in seed_list:
        config['env_name'] = env
        config['seed'] = seed
        print(f"ENV:{env}  AGENT:DQN  SEED:{seed}")
        agent = DQN(config)
        evaluations = agent.run()
        results.append([evaluations, agent.get_score_required_to_win(), seed])
        print('-' * 100)
        print('-' * 100)
    data = {'DQN': results}
    path = Path(f'../results/data/{env}_DQN.json')
    json_str = json.dumps(data, indent=4)
    path.write_text(json_str)
